from PyQt6.QtWidgets import QTextBrowser, QVBoxLayout, QWidget, QScrollArea, QHBoxLayout, QLabel
from PyQt6.QtCore import Qt, pyqtSignal
from PyQt6.QtGui import QTextCursor
import markdown
from styles import (
    USER_BUBBLE_STYLE, USER_MESSAGE_STYLE,
    ASSISTANT_BUBBLE_STYLE, ASSISTANT_MESSAGE_STYLE,
    CODE_BLOCK_CONTAINER_STYLE, CODE_BLOCK_HEADER_STYLE,
    CODE_BLOCK_CONTENT_STYLE, TEXT_DISPLAY_STYLE,
    ASSISTANT_CHUNK_STYLE, ASSISTANT_CHUNK_MESSAGE_STYLE,
    USER_SEPARATOR_STYLE, MESSAGE_SEPARATOR_STYLE,
    ERROR_MESSAGE_STYLE, CODE_STYLE
)

class ChatDisplay(QScrollArea):
    edit_message = pyqtSignal(str, str)

    def __init__(self):
        super().__init__()
        self.init_ui()
        self.messages = []
        self.message_max_width = 250

    def init_ui(self):
        self.layout = QVBoxLayout(self)
        self.text_display = QTextBrowser()
        self.text_display.setOpenExternalLinks(False)
        self.text_display.setReadOnly(True)
        self.text_display.setStyleSheet(TEXT_DISPLAY_STYLE)
        self.layout.addWidget(self.text_display)
        self.layout.setContentsMargins(1, 1, 1, 1)

    def add_message(self, text, role="user"):
        """添加新消息"""
        if role == "user":
            # 用户消息直接显示原文本
            html = self._create_user_message_html(text)
        else:
            # AI 助手的消息需要 markdown 渲染
            text_with_code_blocks = self._format_code_blocks(text)
            html = self._create_assistant_message_html(text_with_code_blocks)
        
        self.messages.append({"role": role, "content": text})
        cursor = self.text_display.textCursor()
        cursor.movePosition(QTextCursor.MoveOperation.End)
        cursor.insertHtml(html)
        self.text_display.verticalScrollBar().setValue(
            self.text_display.verticalScrollBar().maximum()
        )

    def display_assistant_message(self, message, is_chunk=False):
        """显示 AI 消息"""
        if is_chunk:
            if self.messages and self.messages[-1]["role"] == "assistant":
                current_text = self.messages[-1]["content"]
                new_text = current_text + message
                self.messages[-1]["content"] = new_text
                
                # 更新最后一条消息的显示
                cursor = self.text_display.textCursor()
                cursor.movePosition(QTextCursor.MoveOperation.End)
                cursor.insertHtml(f"""
                <div style="{ASSISTANT_CHUNK_STYLE}">
                    <div style="{ASSISTANT_CHUNK_MESSAGE_STYLE}">{message}</div>
                </div>
                """)
            else:
                # 第一个 chunk，创建新消息
                self.messages.append({"role": "assistant", "content": message})
                cursor = self.text_display.textCursor()
                cursor.movePosition(QTextCursor.MoveOperation.End)
                cursor.insertHtml(f"""
                <div style="{ASSISTANT_CHUNK_STYLE}">
                    <div style="{ASSISTANT_CHUNK_MESSAGE_STYLE}">{message}</div>
                </div>
                """)
        else:
            # 完整消息，使用 markdown 渲染
            if self.messages and self.messages[-1]["role"] == "assistant":
                # 更新最后一条消息的内容
                self.messages[-1]["content"] = message
                
                # 清空显示并重新渲染所有消息
                self.text_display.clear()
                for msg in self.messages:
                    if msg["role"] == "user":
                        html = self._create_user_message_html(msg["content"])
                    else:
                        text_with_code_blocks = self._format_code_blocks(msg["content"])
                        html = self._create_assistant_message_html(text_with_code_blocks)
                    cursor = self.text_display.textCursor()
                    cursor.movePosition(QTextCursor.MoveOperation.End)
                    cursor.insertHtml(html)
            else:
                # 如果没有之前的消息，直接添加新消息
                self.add_message(message, role="assistant")
                
        # 确保滚动到最新内容
        self.scroll_to_bottom()

    def update_last_message(self, text):
        """更新最后一条消息的内容"""
        if self.messages:
            self.messages[-1]["content"] = text
            self.text_display.clear()
            for msg in self.messages:
                if msg["role"] == "user":
                    # 确保每条用户消息都有分隔符
                    html = self._create_user_message_html(msg["content"])
                else:
                    text_with_code_blocks = self._format_code_blocks(msg["content"])
                    html = self._create_assistant_message_html(text_with_code_blocks)
                cursor = self.text_display.textCursor()
                cursor.movePosition(QTextCursor.MoveOperation.End)
                cursor.insertHtml(html)
            
            # 确保滚动到最新内容
            self.scroll_to_bottom()

    def scroll_to_bottom(self):
        """滚动到最新内容"""
        scrollbar = self.text_display.verticalScrollBar()
        scrollbar.setValue(scrollbar.maximum())

    def _create_user_message_html(self, text):
        """创建用户消息的 HTML - 右侧绿色气泡"""
        return f"""
        <div style="{USER_BUBBLE_STYLE}">
            <span style="{USER_MESSAGE_STYLE}">{text}</span>
        </div>
        <div style="{USER_SEPARATOR_STYLE}"></div>
        """

    def _create_assistant_message_html(self, text):
        """创建助手消息的 HTML - 左侧白色气泡"""
        text_with_code_blocks = self._format_code_blocks(text)
        html = markdown.markdown(text_with_code_blocks, extensions=[
            'markdown.extensions.fenced_code',
            'markdown.extensions.tables',
            'markdown.extensions.nl2br',
        ])
        
        parts = html.split('<div style="' + CODE_BLOCK_CONTAINER_STYLE.strip() + '">')
        result = []
        
        for i, part in enumerate(parts):
            if i == 0:
                if part.strip():
                    result.append(f"""
                    <div style="{ASSISTANT_BUBBLE_STYLE}">
                        <div style="{ASSISTANT_MESSAGE_STYLE}">{part}</div>
                    </div>
                    <div style="{MESSAGE_SEPARATOR_STYLE}"></div>
                    """)
            else:
                code_part = part.split('</div>', 1)
                if len(code_part) > 1:
                    result.append(f'<div style="{CODE_BLOCK_CONTAINER_STYLE}">' + code_part[0] + '</div>')
                    if code_part[1].strip():
                        result.append(f"""
                        <div style="{ASSISTANT_BUBBLE_STYLE}">
                            <div style="{ASSISTANT_MESSAGE_STYLE}">{code_part[1]}</div>
                        </div>
                        <div style="{MESSAGE_SEPARATOR_STYLE}"></div>
                        """)
        
        return '\n'.join(result)

    def _format_code_blocks(self, text):
        """格式化文本中的代码块"""
        import re
        
        # 更精确的代码块匹配模式，包括语言标识和文件路径
        code_block_pattern = r'```([\w:\/\.-]*)\n((?:(?!```).|\n)+?)```'
        
        def replace_code_block(match):
            lang_or_path = match.group(1) or ''
            code = match.group(2)
            
            # 分离语言和文件路径（如果有的话）
            if ':' in lang_or_path:
                language, path = lang_or_path.split(':', 1)
            else:
                language, path = lang_or_path, ''
            
            return self._create_code_block_html(code.strip(), language, path)
        
        # 处理所有代码块
        formatted = re.sub(code_block_pattern, replace_code_block, text, flags=re.DOTALL)
        return formatted

    def _create_code_block_html(self, code, language='', path=''):
        """创建代码块的 HTML"""
        import html
        escaped_code = html.escape(code)
        
        # 构建标题分
        header = language
        if path:
            header = f"{language}:{path}"
        
        # 修改代码块的容器样式，确保它是独立的块
        return f'''
        <div style="{CODE_BLOCK_CONTAINER_STYLE}">
            <div style="{CODE_BLOCK_HEADER_STYLE}">{header}</div>
            <pre style="{CODE_BLOCK_CONTENT_STYLE}">
                <code style="{CODE_STYLE}">{escaped_code}</code>
            </pre>
        </div>
        <div style="height: 10px;"></div>  <!-- 添加空行确保代码块与其他内容分开 -->
        '''

    def display_error(self, error_message):
        """显示错误消息"""
        error_html = f"""
        <div style="margin: 10px 0; text-align: center;">
            <div style="{ERROR_MESSAGE_STYLE}">
                ⚠️ {error_message}
            </div>
        </div>
        """
        cursor = self.text_display.textCursor()
        cursor.movePosition(QTextCursor.MoveOperation.End)
        cursor.insertHtml(error_html)

    def clear_messages(self):
        """清空所有消息"""
        self.messages.clear()
        self.text_display.clear()

    def get_messages(self):
        """获取所有消息历史"""
        return self.messages

    def create_message_widget(self, content, role="assistant"):
        message_widget = QWidget()
        layout = QHBoxLayout(message_widget)
        
        text_label = QLabel(content)
        text_label.setWordWrap(True)
        text_label.setMaximumWidth(self.message_max_width)